# db.py
import sqlite3
import os

import numpy as np


class EmployeeDB:
    """员工数据库操作类"""
    def __init__(self, db_path='employees.db'):
        self.db_path = db_path
        self._create_table()

    def _create_table(self):
        """创建员工数据表"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute('''
             CREATE TABLE IF NOT EXISTS employees (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                employee_id TEXT UNIQUE NOT NULL,
                name TEXT NOT NULL,
                avatar_path TEXT NOT NULL,
                face_encoding BLOB 
            )
            ''')
            conn.commit()

    # db.py
    def add_employee(self, employee_id, name, avatar_path, face_encoding=None):
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO employees (employee_id, name, avatar_path, face_encoding)
                VALUES (?, ?, ?, ?)
            ''', (employee_id, name, avatar_path, face_encoding))
            conn.commit()
            return cursor.lastrowid

    def update_employee(self, employee_id, name, avatar_path, face_encoding=None):
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute('''
                UPDATE employees
                SET name = ?, avatar_path = ?, face_encoding = ?
                WHERE employee_id = ?
            ''', (name, avatar_path, face_encoding, employee_id))
            conn.commit()
            return cursor.rowcount > 0

    def delete_employee(self, employee_id):
        """删除员工"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute('DELETE FROM employees WHERE employee_id = ?', (employee_id,))
            conn.commit()
            return cursor.rowcount > 0

    # db.py
    def get_all_employees(self):
        """获取所有员工（包括人脸向量）"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT id, employee_id, name, avatar_path, face_encoding FROM employees')
            employees = []
            for row in cursor.fetchall():
                # 将二进制 blob 转换为 numpy 数组（如果存在）
                face_encoding = None
                if row[4]:  # face_encoding 字段
                    try:
                        # 从二进制数据重建 numpy 数组
                        face_encoding = np.frombuffer(row[4], dtype=np.float64)
                    except Exception as e:
                        print(f"转换人脸向量时出错: {e}")
                        face_encoding = None

                employees.append({
                    "id": row[0],
                    "employee_id": row[1],
                    "name": row[2],
                    "avatar": row[3],
                    "face_encoding": face_encoding  # 添加人脸向量
                })
            return employees

    def get_employee_by_id(self, employee_id):
        """根据工号获取员工（包括人脸向量）"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT id, employee_id, name, avatar_path, face_encoding
                FROM employees
                WHERE employee_id = ?
            ''', (employee_id,))
            row = cursor.fetchone()
            if row:
                # 将二进制 blob 转换为 numpy 数组（如果存在）
                face_encoding = None
                if row[4]:  # face_encoding 字段
                    try:
                        # 从二进制数据重建 numpy 数组
                        face_encoding = np.frombuffer(row[4], dtype=np.float64)
                    except Exception as e:
                        print(f"转换人脸向量时出错: {e}")
                        face_encoding = None

                return {
                    "id": row[0],
                    "employee_id": row[1],
                    "name": row[2],
                    "avatar": row[3],
                    "face_encoding": face_encoding  # 添加人脸向量
                }
            return None